import networkx as nx
import torch
import numpy as np
from typing import Dict, Any, List, Optional, Tuple, Union
from app.services.dhgam.thgrl_service import TemporalHeterogeneousGraphRepresentationLearning
from app.schemas.dhgam.thgrl_schema import (
    AlgorithmResponse, AlgorithmMiddleResponse, InputParams, OutputParams,
    THGRLOutputParams, NodeEmbedding, LinkPrediction, TaskStatus
)
from app.schemas.metrics_schema import LineChart, ScatterPlot
from app.utils.logger import logger

class THGRLService:
    def __init__(self):
        """初始化THGRL服务"""
        pass
        
    async def run_thgrl(self, input_params: Dict[str, Any], task_id: str) -> Union[AlgorithmResponse, AlgorithmMiddleResponse]:
        """
        运行THGRL算法
        
        Args:
            input_params: 算法输入参数
            task_id: 任务ID
            
        Returns:
            Union[AlgorithmResponse, AlgorithmMiddleResponse]: 算法结果
        """
        try:
            thgrl_params = input_params.get("thgrl_params", {})
            
            # 从输入参数中获取图序列数据
            graph_sequence_data = thgrl_params.get("graph_sequence", [])
            embedding_dim = thgrl_params.get("embedding_dim", 64)
            time_steps = thgrl_params.get("time_steps", 5)
            learning_rate = thgrl_params.get("learning_rate", 0.01)
            epochs = thgrl_params.get("epochs", 100)
            
            # 修复: 确保节点ID和边的源/目标ID在处理前都转换为Python原生字符串
            for time_step, graph in enumerate(graph_sequence_data):
                nodes = graph.get("nodes", [])
                edges = graph.get("edges", [])
                
                # 确保节点ID是字符串类型
                for node in nodes:
                    if "id" in node and not isinstance(node["id"], str):
                        node["id"] = str(node["id"])
                
                # 确保边的源和目标是字符串类型
                for edge in edges:
                    if "source" in edge and not isinstance(edge["source"], str):
                        edge["source"] = str(edge["source"])
                    if "target" in edge and not isinstance(edge["target"], str):
                        edge["target"] = str(edge["target"])
            
            # 转换为NetworkX图序列
            graph_sequence = self._convert_to_nx_graphs(graph_sequence_data)
            
            # 初始化并训练THGRL模型
            thgrl = TemporalHeterogeneousGraphRepresentationLearning(
                embedding_dim=embedding_dim,
                time_steps=time_steps,
                learning_rate=learning_rate,
                epochs=epochs
            )
            
            # 训练模型并获取节点嵌入
            node_embeddings = thgrl.fit(graph_sequence)
            
            # 在处理结果时，确保节点ID也是字符串类型
            node_embeddings_list = []
            for node_id, embedding in node_embeddings.items():
                if not isinstance(node_id, str):
                    node_id = str(node_id)
                
                node_embedding = NodeEmbedding(
                    node_id=node_id,
                    embedding=embedding.tolist(),  # 确保numpy数组转换为Python列表
                    node_type="unknown"  # 这里可能是问题所在，确保node_id是str类型
                )
                node_embeddings_list.append(node_embedding)
            
            # 为最后几个图生成链接预测
            link_predictions = self._generate_link_predictions(thgrl, graph_sequence[-1])
            
            # 生成可视化指标
            # metrics = self._generate_metrics(node_embeddings, graph_sequence)
            metrics = []
            # 准备输出参数
            thgrl_output = THGRLOutputParams(
                node_embeddings=node_embeddings_list,
                link_predictions=link_predictions,
                algorithm="THGRL",
                parameters={
                    "embedding_dim": embedding_dim,
                    "time_steps": time_steps,
                    "learning_rate": learning_rate,
                    "epochs": epochs
                },
                training_loss=[]  # 在实际实现中应从训练过程中收集
            )
            
            output_params = OutputParams(thgrl_results=thgrl_output)
            
            # 构建最终响应
            result = AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.COMPLETED,
                task_progress=100,
                output_params=output_params,
                metrics=metrics
            )
            
            return result
            
        except Exception as e:
            logger.error(f"THGRL算法执行失败: {str(e)}")
            
            # 返回错误响应
            return AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.FAILED,
                task_progress=0,
                error_message=str(e),
                output_params=OutputParams(),
                metrics=[]
            )
    
    def _convert_to_nx_graphs(self, graph_sequence_data: List[Dict[str, Any]]) -> List[nx.Graph]:
        """
        将JSON格式的图序列数据转换为NetworkX图对象序列
        
        Args:
            graph_sequence_data: JSON格式的图序列数据
            
        Returns:
            List[nx.Graph]: NetworkX图对象序列
        """
        graph_sequence = []
        
        for graph_data in graph_sequence_data:
            G = nx.Graph()
            
            # 添加节点和节点属性
            for node_data in graph_data.get("nodes", []):
                node_id = node_data.get("id")
                node_type = node_data.get("type", "unknown")
                node_attrs = {k: v for k, v in node_data.items() if k not in ["id", "type"]}
                node_attrs["type"] = node_type
                
                G.add_node(node_id, **node_attrs)
            
            # 添加边和边属性
            for edge_data in graph_data.get("edges", []):
                source = edge_data.get("source")
                target = edge_data.get("target")
                edge_type = edge_data.get("type", "default")
                edge_attrs = {k: v for k, v in edge_data.items() if k not in ["source", "target", "type"]}
                edge_attrs["type"] = edge_type
                
                G.add_edge(source, target, **edge_attrs)
            
            graph_sequence.append(G)
        
        return graph_sequence
    
    def _convert_embeddings(self, node_embeddings: Dict[Any, np.ndarray]) -> List[NodeEmbedding]:
        """
        将节点嵌入字典转换为NodeEmbedding对象列表
        
        Args:
            node_embeddings: 节点嵌入字典，键为节点ID，值为嵌入向量
            
        Returns:
            List[NodeEmbedding]: 节点嵌入对象列表
        """
        result = []
        
        for node_id, embedding in node_embeddings.items():
            node_type = "unknown"
            if hasattr(node_id, "get") and callable(getattr(node_id, "get")):
                node_type = node_id.get("type", "unknown")
                node_id = str(node_id.get("id", node_id))
            
            result.append(NodeEmbedding(
                node_id=str(node_id),
                embedding=embedding.tolist(),
                node_type=node_type
            ))
        
        return result
    
    def _generate_link_predictions(self, thgrl: TemporalHeterogeneousGraphRepresentationLearning, graph: nx.Graph) -> List[LinkPrediction]:
        """
        生成链接预测结果
        
        Args:
            thgrl: THGRL模型
            graph: 测试图
            
        Returns:
            List[LinkPrediction]: 链接预测结果列表
        """
        predictions = []
        nodes = list(graph.nodes())
        
        # 从图中随机选择一些节点对进行链接预测
        for _ in range(min(100, len(nodes) * 10)):  # 限制预测数量
            source = np.random.choice(nodes)
            target = np.random.choice(nodes)
            
            if source != target:
                # 修复：确保source和target是Python原生字符串类型
                if isinstance(source, np.str_) or not isinstance(source, str):
                    source = str(source)
                if isinstance(target, np.str_) or not isinstance(target, str):
                    target = str(target)
                    
                edge_data = graph.get_edge_data(source, target) or {}
                edge_type = edge_data.get("type", "default") if edge_data else "default"
                
                # 使用THGRL模型预测链接概率
                prob = thgrl.predict_link(source, target, edge_type=edge_type)
                
                predictions.append(LinkPrediction(
                    source_node=str(source),
                    target_node=str(target),
                    probability=float(prob),
                    edge_type=edge_type
                ))
        
        # 按概率降序排序
        predictions.sort(key=lambda x: x.probability, reverse=True)
        
        return predictions[:50]  # 返回前50个最可能的链接
    
    def _generate_metrics(self, node_embeddings: Dict[Any, np.ndarray], graph_sequence: List[nx.Graph]) -> List[Any]:
        """
        生成可视化指标
        
        Args:
            node_embeddings: 节点嵌入字典
            graph_sequence: 图序列
            
        Returns:
            List[Any]: 可视化指标列表
        """
        metrics = []
        
        # 1. 生成嵌入空间的2D散点图 (使用t-SNE降维)
        try:
            from sklearn.manifold import TSNE
            import numpy as np
            
            embeddings = np.array(list(node_embeddings.values()))
            
            if len(embeddings) > 5:  # 确保有足够的数据点
                # 降维到2D
                tsne = TSNE(n_components=2, random_state=42)
                embeddings_2d = tsne.fit_transform(embeddings)
                
                # 准备散点图数据
                scatter_data = []
                for i, node_id in enumerate(node_embeddings.keys()):
                    # 修复：针对node_id的类型检查和处理
                    str_node_id = str(node_id)  # 确保转换为字符串
                    
                    # 获取节点类型，不再尝试调用numpy.str_对象的get方法
                    node_type = "unknown"
                    if isinstance(node_id, dict):  # 只在node_id是字典时才尝试获取type
                        node_type = node_id.get("type", "unknown")
                    elif graph_sequence and str_node_id in graph_sequence[-1]:
                        # 尝试从最后一个图中获取节点类型
                        node_attrs = graph_sequence[-1].nodes.get(str_node_id, {})
                        node_type = node_attrs.get("type", "unknown")
                    
                    scatter_data.append({
                        "x": float(embeddings_2d[i, 0]),
                        "y": float(embeddings_2d[i, 1]),
                        "id": str_node_id,
                        "type": node_type
                    })
                
                metrics.append(ScatterPlot(
                    title="节点嵌入2D可视化",
                    xLabel="t-SNE维度1",
                    yLabel="t-SNE维度2",
                    data=scatter_data
                ))
        except Exception as e:
            logger.warning(f"生成嵌入散点图失败: {str(e)}")
        
        # 2. 生成图统计随时间变化的折线图
        try:
            time_points = list(range(len(graph_sequence)))
            
            # 节点数量随时间变化
            node_counts = [len(g.nodes()) for g in graph_sequence]
            
            node_count_data = [{"x": t, "y": count} for t, count in zip(time_points, node_counts)]
            
            metrics.append(LineChart(
                title="节点数量随时间变化",
                xLabel="时间步",
                yLabel="节点数量",
                data=node_count_data
            ))
            
            # 边数量随时间变化
            edge_counts = [len(g.edges()) for g in graph_sequence]
            
            edge_count_data = [{"x": t, "y": count} for t, count in zip(time_points, edge_counts)]
            
            metrics.append(LineChart(
                title="边数量随时间变化",
                xLabel="时间步",
                yLabel="边数量",
                data=edge_count_data
            ))
        except Exception as e:
            logger.warning(f"生成图统计折线图失败: {str(e)}")
        
        return metrics 